package cn.hinglo.common.ratelimit;

import cn.hinglo.common.ratelimit.config.LimitCacheTypeEnum;
import cn.hinglo.common.ratelimit.config.LimitRenderType;
import cn.hinglo.common.ratelimit.config.LimitationInfo;
import cn.hinglo.common.utils.ResultUtils;
import cn.hutool.core.util.NumberUtil;
import cn.hutool.core.util.StrUtil;
import cn.hutool.db.nosql.redis.RedisDS;
import cn.hutool.setting.Setting;
import com.google.common.cache.CacheBuilder;
import com.google.common.cache.CacheLoader;
import com.google.common.cache.LoadingCache;
import com.google.common.util.concurrent.RateLimiter;
import redis.clients.jedis.Jedis;

import javax.servlet.ServletException;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import java.io.IOException;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.TimeUnit;


/**
 * @author HingLo
 * @createTime: 2018/9/25 10:03
 * @description: 限流管理工具
 */
public final class RateLimitManager {

    /**
     * 自身单例对象
     */
    private static final RateLimitManager me = new RateLimitManager();
    /**
     * redis数据源
     */
    private static final RedisDS redisDs;
    /**
     * 限流缓存类型,目前支持(redis与内存两种方式)
     */
    private final static LimitCacheTypeEnum requestType;

    static {
        Setting setting = new Setting("redis.setting");
        if (!setting.isEmpty()) {
            redisDs = RedisDS.create(setting, null);
            requestType = LimitCacheTypeEnum.REDIS;
        } else {

            redisDs = null;
            requestType = LimitCacheTypeEnum.MEMORY;
        }
    }

    /**
     * 创建一个线程安全的map类型的令牌桶,使用内存缓存限流信息
     */
    private LoadingCache<String, RateLimiter> requestRateLimiterMap = CacheBuilder.newBuilder().expireAfterAccess(60, TimeUnit.SECONDS)
            .build(new CacheLoader<String, RateLimiter>() {

                @Override
                public RateLimiter load(String key) {
                    return null;
                }
            });

    /**
     * 构造函数,默认要初始化配置信息
     */
    private RateLimitManager() {
    }

    /**
     * 单例模式
     *
     * @return 返回单例对象
     */
    public static RateLimitManager me() {
        return me;
    }

    /**
     * 初始化一个限流的目标
     *
     * @param target 目标对象,ip,url,user,request等
     * @param rate   每秒通过的路由
     * @return 返回限流对象
     */
    private RateLimiter initRateLimiter(String target, double rate) {
        // 保留2位小数就好
        rate = NumberUtil.round(rate, 2).doubleValue();

        RateLimiter limiter = getLimiter(target);
        if (limiter == null) {
            limiter = RateLimiter.create(rate);
            setLimiter(target, limiter);
            return limiter;
        }

        if (limiter.getRate() == rate) {
            return limiter;
        }
        limiter.setRate(rate);
        setLimiter(target, limiter);
        return limiter;
    }

    /**
     * 通过指定的可以来返回限流对象
     *
     * @param target 目标对象
     * @return 返回限流对象, 默认使用内存作为缓存
     */
    private RateLimiter getLimiter(String target) {
        try {
            return requestRateLimiterMap.get(target);
        } catch (ExecutionException | RuntimeException e) {
            return null;
        }
    }

    /**
     * 添加或者更新限流信息
     *
     * @param target  目标
     * @param limiter 限流对象
     */
    private void setLimiter(String target, RateLimiter limiter) {
        requestRateLimiterMap.put(target, limiter);
    }

    /**
     * 限流方法
     *
     * @param target 目标
     * @param rate   限流次数
     * @param second 限流时间
     * @return 返回结果
     */
    private boolean limit(String target, int rate, int second) {
        // 使用内存方案
        if (requestType == LimitCacheTypeEnum.MEMORY) {

            RateLimiter rateLimiter = initRateLimiter(target, rate * 1.0 / second);
            return rateLimiter.tryAcquire();

        } else {
            // 下面是使用redis方式限流
            String val = redisGetLimit(target);
            if (val == null) {
                redisSetLimit(target, 1, second);
                return true;
            } else {
                int count = Integer.parseInt(val);
                if (count < rate) {
                    // 请求次数+1
                    redisSetLimit(target, count + 1, -1);
                    return true;
                }
                return false;
            }
        }
    }

    /**
     * redis限流方式的存入
     *
     * @param target 目标
     * @param rate   次数
     * @param second 制定的时间
     */
    private void redisSetLimit(String target, int rate, int second) {
        Jedis jedis = redisDs.getJedis();
        // 如果second=-1 表示是更新
        if (second == -1) {
            jedis.set(target, String.valueOf(rate));
            // 剩余时间
            long millis = jedis.ttl(target);
            // 更新有效期的时间
            jedis.pexpire(target, millis);
        } else {
            jedis.setex(target, second, String.valueOf(rate));
        }

        jedis.close();
    }

    /**
     * 通过目标 返回限流信息
     *
     * @param target 目标对象
     * @return 返回结果
     */
    private String redisGetLimit(String target) {
        Jedis jedis = redisDs.getJedis();
        String value = jedis.get(target);
        jedis.close();
        return value;
    }

    /**
     * 并发限流
     *
     * @param target         请求的目标
     * @param limitationInfo 限流信息
     * @param response       响应信息
     * @return 返回限流结果
     */
    public boolean rateLimit(String target, LimitationInfo limitationInfo, HttpServletResponse response, HttpServletRequest request)
            throws IOException, ServletException {
        boolean res = limit(target, limitationInfo.getRate(), limitationInfo.getSecond());
        if (!res) {
            this.render(limitationInfo.getRenderType(), limitationInfo.getRenderContent(), response, request);
        }
        return res;
    }

    /**
     * 根据类型返回结果
     *
     * @param renderType    返回类型
     * @param renderContent 返回的内容
     * @param response      响应对象
     * @param request       请求对象
     * @throws IOException      IO
     * @throws ServletException servlet
     */
    private void render(String renderType, String renderContent, HttpServletResponse response, HttpServletRequest request)
            throws IOException, ServletException {
        if (StrUtil.equalsIgnoreCase(renderType, LimitRenderType.JSON)) {
            // 如果是返回json格式的数据
            ResultUtils.resp(response, ResultUtils.error(renderContent));
        } else if (renderType.equalsIgnoreCase(LimitRenderType.PAGE)) {
            // 返回页面
            request.getRequestDispatcher(renderContent).forward(request, response);
        } else if (StrUtil.equalsIgnoreCase(renderType, LimitRenderType.REDIRECT)) {
            // 重定向处理
            response.sendRedirect(renderContent);
        }
    }

}
